from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from sqlalchemy import Column, Integer, String, MetaData
from sqlalchemy.ext.declarative import declarative_base
import pandas as pd
 
# 定义数据库连接URL
DATABASE_URL = "sqlite+aiosqlite:///yourdatabase.db"  # 你可以使用其他数据库URL，如PostgreSQL的"postgresql+asyncpg://user:password@localhost/dbname"
 
# 创建异步引擎
engine = create_async_engine(DATABASE_URL, echo=True, future=True)
 
# 创建异步会话类
async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
 
Base = declarative_base()
 
class MyModel(Base):
    __tablename__ = 'my_table'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    value = Column(Integer)
 
# 创建表（如果需要）
async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
 
# 异步写入数据到数据库
async def write_data_to_db(df):
    async with async_session() as session:
        for index, row in df.iterrows():
            new_record = MyModel(name=row['name'], value=row['value'])
            session.add(new_record)
        await session.commit()
 
# 在你的FastAPI路由中使用它
from fastapi import FastAPI
app = FastAPI()
 
@app.on_event("startup")
async def startup_event():
    await init_db()
 
@app.post("/write_data/")
async def write_data(data: dict):
    df = pd.DataFrame([data])  # 假设你从请求中获取数据并转换为DataFrame
    await write_data_to_db(df)
    return {"message": "Data written successfully"}